import os
import torch.utils.data
import cv2
import numpy as np
from libs.utils import  show_sample_withplt, show_sample_withocv, getFilepath, np2Tensor
from pathlib import *

def get_image_and_mask(rootdir):
    rootpath = Path(rootdir)
    imagedir = rootpath/"image"
    labeldir = rootpath/"mask"

    image_list = []
    mask_list = []

    for maskpath in labeldir.iterdir():
        imagepartname = maskpath.stem
        imagepartdir = imagedir/imagepartname
        for imgpath in imagepartdir.glob("*"):
                image_list.append(imgpath)
                mask_list.append(maskpath)
    return image_list,mask_list


class HairDataset(torch.utils.data.Dataset):
    def __init__(self, data_folder, image_size=(480, 480), mask_size=(480, 480)):
        self.data_folder = data_folder
        if not os.path.exists(self.data_folder):
            raise Exception("%s  not exists." % self.data_folder)

        self.image_list, self.mask_list = get_image_and_mask(data_folder)
        self.image_size = image_size
        self.mask_size = mask_size

    def image_preprocess(self, image):
        image = cv2.resize(image, self.image_size)
        image = image.astype(np.float32)
        image = ((image / 255.0) - 0.5) / 0.5
        image = np2Tensor(image)
        return image


    def mask_preprocess(self, mask):
        mask = cv2.resize(mask, self.mask_size)
        testimg = (mask[:,:,0]==255).astype(np.uint8) *  (mask[:,:,1]==0).astype(np.uint8) * (mask[:,:,2]==0).astype(np.uint8)

        print(testimg.shape)
        cv2.imshow("f", 255*testimg)

        cv2.waitKey(0)

        combinemask = np.zeros(self.mask_size, dtype=np.uint8)
        #combinemask = combinemask+r_1
        #
        # mask = mask.astype(np.float32)
        # mask = (mask / 255.0)
        combinemask = np2Tensor(combinemask)
        return combinemask


    def __getitem__(self, index):
        image = cv2.imread(str(self.image_list[index]))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(str(self.mask_list[index]) )
        # mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)
        image = self.image_preprocess(image)
        mask = self.mask_preprocess(mask)
        return image, mask

    def __len__(self):
        return len(self.image_list)




if __name__ == "__main__":

    # input_shape = (1920, 1080)
    # output_shape = (1920, 1080)
    input_shape = (800, 800)
    output_shape = (800, 800)
    for image, mask in HairDataset('/home/hanson/work/roadSegmentation/dataset',
                              image_size=input_shape,
                              mask_size=output_shape):

        image = image.cpu().numpy()
        image = np.transpose(image, (1, 2, 0))
        image = 255 * (image * 0.5 + 0.5)
        image = image.astype(np.uint8)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        mask = mask.cpu().numpy()[0]

        mask = mask.astype(np.uint8)

        # cv2.imshow("s", image)
        # cv2.imshow("m", mask)
        # cv2.waitKey(0)
